package cn.bincker.mybatis.encrypt.data.migration.impl;

import cn.bincker.mybatis.encrypt.core.EncryptExecutor;
import cn.bincker.mybatis.encrypt.data.migration.entity.ColumnInfo;
import cn.bincker.mybatis.encrypt.data.migration.entity.MigrateTable;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.io.Closeable;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
public class EncryptTask implements Runnable, Closeable {
    private final EncryptExecutor encryptExecutor;
    @Getter
    private boolean running;
    private boolean destroyed = false;
    private final BlockingQueue<Object[]> queue;
    private final BlockingQueue<Object[]> nextQueue;
    private MigrateTable migrateTable;
    private final AtomicLong encryptRowCount;
    private final AtomicLong encryptFailRowCount;
    private final Consumer<EncryptTask> onEncryptThreadInitListener;
    private final Consumer<EncryptTask> onEncryptThreadDestroyListener;
    private final boolean failContinue;
    private final int nextQueueLimit;
    private final Object busyLock = new Object();
    private transient boolean busy;

    public EncryptTask(
            EncryptExecutor encryptExecutor,
            BlockingQueue<Object[]> queue,
            BlockingQueue<Object[]> nextQueue,
            AtomicLong encryptRowCount, AtomicLong encryptFailRowCount, Consumer<EncryptTask> onEncryptThreadInitListener,
            Consumer<EncryptTask> onEncryptThreadDestroyListener,
            boolean failContinue,
            int nextQueueLimit
    ) {
        this.encryptRowCount = encryptRowCount;
        this.encryptFailRowCount = encryptFailRowCount;
        this.nextQueueLimit = nextQueueLimit;
        this.running = true;
        this.encryptExecutor = encryptExecutor;
        this.queue = queue;
        this.nextQueue = nextQueue;
        this.onEncryptThreadInitListener = onEncryptThreadInitListener;
        this.onEncryptThreadDestroyListener = onEncryptThreadDestroyListener;
        this.failContinue = failContinue;
    }

    public void setMigrateTable(MigrateTable table){
        this.migrateTable = table;
    }

    @Override
    public void run() {
        running = true;
        destroyed = false;
        if (onEncryptThreadInitListener != null) onEncryptThreadInitListener.accept(this);
        try {
            mainLoop:
            while (running) {
                Object[] row;
                try {
                    row = queue.poll(1, TimeUnit.SECONDS);
                    if (row == null) {
                        synchronized (busyLock){
                            busy = false;
                            busyLock.notifyAll();
                        }
                        continue ;
                    }
                } catch (InterruptedException ignore) {
                    continue;
                }
                busy = true;
                for (ColumnInfo column : migrateTable.getEncryptColumns()) {
                    var retryTimes = 0;
                    while (true) {
                        try {
                            row[column.getIndex()] = encryptExecutor.encrypt(column.getEncryptProperty(), null, row[column.getIndex()]);
                            break;
                        } catch (Throwable e) {
                            if (retryTimes ++ > 3) {
                                encryptFailRowCount.incrementAndGet();
                                if (failContinue) {
                                    log.error("encrypt row fail: column={}\tkeys={}", migrateTable.getSrcTableName() + "." + column.getName(), migrateTable.getKeys().stream().map(k -> row[k.getIndex()].toString()).collect(Collectors.joining(",")), e);
                                    continue mainLoop;
                                } else {
                                    throw e;
                                }
                            }
                        }
                    }
                }
                waitingForSpend();
                nextQueue.add(row);
                encryptRowCount.incrementAndGet();
                if (queue.isEmpty()) {
                    synchronized (queue) {
                        queue.notifyAll();
                    }
                }
            }
        }catch (Throwable e){
            log.error("encrypt thread error", e);
            running = false;
        }finally {
            if(onEncryptThreadDestroyListener != null) onEncryptThreadDestroyListener.accept(this);
        }
        destroyed = true;
        synchronized (this){
            this.notifyAll();
        }
    }

    private void waitingForSpend() throws InterruptedException {
        if (nextQueue.size() < nextQueueLimit) return;
        synchronized (nextQueue){
            while (nextQueue.size() > nextQueueLimit)
                nextQueue.wait(1000L);
        }
    }

    @Override
    public void close() {
        running = false;
    }

    public void waitDestroyed() {
        if (destroyed) return;
        synchronized (this){
            try {
                this.wait();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public void waitBusy(){
        while (busy){
            synchronized (busyLock){
                if (!busy) return;
                try {
                    busyLock.wait(1000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }
}
